// src/client.rs
use tokio::net::{TcpListener, TcpStream};
use crate::transport::{Transport, WebSocketTransport, TcpTransport};
use tokio::io::{AsyncReadExt, AsyncWriteExt, AsyncBufReadExt};
use url::Url;
use rustls;
use rustls::pki_types;

#[derive(Debug)]
struct NoopServerCertVerifier;

impl rustls::client::danger::ServerCertVerifier for NoopServerCertVerifier {
    fn verify_server_cert(
        &self,
        _end_entity: &pki_types::CertificateDer<'_>,
        _intermediates: &[pki_types::CertificateDer<'_>],
        _server_name: &pki_types::ServerName<'_>,
        _ocsp_response: &[u8],
        _now: pki_types::UnixTime,
    ) -> Result<rustls::client::danger::ServerCertVerified, rustls::Error> {
        Ok(rustls::client::danger::ServerCertVerified::assertion())
    }

    fn verify_tls12_signature(
        &self,
        _message: &[u8],
        _cert: &pki_types::CertificateDer<'_>,
        _dss: &rustls::DigitallySignedStruct,
    ) -> Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
        Ok(rustls::client::danger::HandshakeSignatureValid::assertion())
    }

    fn verify_tls13_signature(
        &self,
        _message: &[u8],
        _cert: &pki_types::CertificateDer<'_>,
        _dss: &rustls::DigitallySignedStruct,
    ) -> Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
        Ok(rustls::client::danger::HandshakeSignatureValid::assertion())
    }

    fn supported_verify_schemes(&self) -> Vec<rustls::SignatureScheme> {
        rustls::crypto::ring::default_provider().signature_verification_algorithms.supported_schemes()
    }
}

pub mod protocol {
    use tokio::io::{AsyncReadExt, AsyncWriteExt};
    use tokio::net::TcpStream;
    use crate::transport::Transport;

    pub async fn forward_streams(
        mut tcp: TcpStream,
        mut transport: Box<dyn Transport + Send>,
    ) -> anyhow::Result<()> {
        let mut buffer = [0u8; 8192];
        loop {
            tokio::select! {
                result = tcp.read(&mut buffer) => {
                    match result {
                        Ok(0) => break,
                        Ok(n) => {
                            if transport.write(&buffer[..n]).await.is_err() {
                                break;
                            }
                        }
                        Err(_) => break,
                    }
                }
                result = transport.read() => {
                    match result {
                        Ok(Some(data)) => {
                            if tcp.write_all(&data).await.is_err() {
                                break;
                            }
                        }
                        Ok(None) => break,
                        Err(_) => break,
                    }
                }
            }
            if transport.is_closed() {
                break;
            }
        }
        Ok(())
    }
}

#[derive(Debug)]
pub struct Config {
    pub listen_addr: String,
    pub server_addr: String,
    pub transport: String,
    pub no_tls: bool,
}

pub async fn start_client(config: Config, mut shutdown_receiver: tokio::sync::oneshot::Receiver<()>, addr_sender: tokio::sync::oneshot::Sender<std::net::SocketAddr>) -> anyhow::Result<()> {
    let listener = TcpListener::bind(&config.listen_addr).await?;
    let local_addr = listener.local_addr()?;
    log::info!("Proxy client listening on {} with {} transport ", local_addr, config.transport);
    let _ = addr_sender.send(local_addr);

    loop {
        tokio::select! {
            _ = &mut shutdown_receiver => {
                log::info!("Client shutting down.");
                break;
            }
            conn = listener.accept() => {
                let (tcp_stream, _) = conn?;
                let server_addr = config.server_addr.clone();
                let no_tls = config.no_tls;

                let transport_type = config.transport.clone();
                tokio::spawn(async move {
                    let result = async {
                        let (tcp_stream, host, port) = handle_proxy_request(tcp_stream).await?;

                        let transport: Box<dyn Transport + Send> = match transport_type.as_str() {
                            "ws" => {
                                if no_tls {
                                    let ws_url = format!("ws://{{}}", server_addr);
                                    let (ws_stream, _) = tokio_tungstenite::connect_async(&ws_url).await?;
                                    log::info!("Connected to WebSocket server: {{}}", ws_url);
                                    let mut ws_transport = WebSocketTransport::new(ws_stream);
                                    ws_transport.send_metadata(&host, port).await?;
                                    Box::new(ws_transport)
                                } else {
                                    let ws_url = format!("wss://{{}}", server_addr);
                                    let url = Url::parse(&ws_url)?;
                                    let host_str = url.host_str().ok_or_else(|| anyhow::anyhow!("Invalid host in URL"))?;
                                    
                                    let config = rustls::ClientConfig::builder_with_provider(
                                        rustls::crypto::ring::default_provider().into()
                                    )
                                    .with_protocol_versions(&[&rustls::version::TLS13])
                                    .unwrap()
                                    .dangerous()
                                    .with_custom_certificate_verifier(std::sync::Arc::new(NoopServerCertVerifier))
                                    .with_no_client_auth();
                                    
                                    let tls_connector = tokio_rustls::TlsConnector::from(std::sync::Arc::new(config));

                                    let stream = TcpStream::connect(&server_addr).await?;
                                    let domain = rustls::pki_types::ServerName::try_from(host_str).unwrap().to_owned();
                                    let stream = tls_connector.connect(domain, stream).await?;

                                    let (ws_stream, _) = tokio_tungstenite::client_async(&ws_url, stream).await?;
                                    log::info!("Connected to WebSocket server: {{}}", server_addr);
                                    let mut ws_transport = WebSocketTransport::new(ws_stream);
                                    ws_transport.send_metadata(&host, port).await?;
                                    Box::new(ws_transport)
                                }
                            },
                            "tcp" => {
                                let tcp_stream = TcpStream::connect(&server_addr).await?;
                                log::info!("Connected to TCP server: {{}}", server_addr);
                                let mut transport = TcpTransport::new(tcp_stream);
                                let metadata = serde_json::json!({
                                    "host": host,
                                    "port": port
                                });
                                let metadata_bytes = metadata.to_string().into_bytes();
                                let len = metadata_bytes.len() as u32;
                                let mut buf = len.to_be_bytes().to_vec();
                                buf.extend_from_slice(&metadata_bytes);
                                transport.write(&buf).await?;
                                Box::new(transport)
                            },
                            _ => return Err(anyhow::anyhow!("Unsupported transport: {{}}", transport_type)),
                        };

                        protocol::forward_streams(tcp_stream, transport).await
                    };

                    if let Err(e) = result.await {
                        log::error!("Error processing connection: {}", e);
                    }
                });
            }
        }
    }
    Ok(())
}

async fn handle_proxy_request(stream: TcpStream) -> anyhow::Result<(TcpStream, String, u16)> {
    let mut buf = [0u8; 1];
    stream.peek(&mut buf).await?;
    if buf[0] == 0x05 {
        log::info!("Detected SOCKS5 proxy request");
        handle_socks5_handshake(stream).await
    } else {
        log::info!("Detected HTTP proxy request");
        handle_http_proxy_request(stream).await
    }
}

async fn handle_http_proxy_request(stream: TcpStream) -> anyhow::Result<(TcpStream, String, u16)> {
    let mut reader = tokio::io::BufReader::new(stream);
    let mut line = String::new();
    reader.read_line(&mut line).await?;

    let parts: Vec<&str> = line.split_whitespace().collect();
    if parts.len() != 3 || parts[0] != "CONNECT" {
        return Err(anyhow::anyhow!("Invalid HTTP proxy request"));
    }

    let host_port: Vec<&str> = parts[1].split(':').collect();
    if host_port.len() != 2 {
        return Err(anyhow::anyhow!("Invalid host:port in CONNECT request"));
    }
    let host = host_port[0].to_string();
    let port = host_port[1].parse::<u16>()?;

    // Read the rest of the headers
    loop {
        line.clear();
        let bytes_read = reader.read_line(&mut line).await?;
        if bytes_read == 0 || line == "\r\n" {
            break;
        }
    }

    let mut stream = reader.into_inner();
    stream.write_all(b"HTTP/1.1 200 Connection established\r\n\r\n").await?;

    Ok((stream, host, port))
}

async fn handle_socks5_handshake(mut stream: TcpStream) -> anyhow::Result<(TcpStream, String, u16)> {
    let mut buf = [0u8; 2];
    stream.read_exact(&mut buf).await?;
    if buf[0] != 0x05 {
        return Err(anyhow::anyhow!("Invalid SOCKS5 version"));
    }
    let nmethods = buf[1];
    let mut methods = vec![0u8; nmethods as usize];
    stream.read_exact(&mut methods).await?;

    if !methods.contains(&0x00) {
        stream.write_all(&[0x05, 0xFF]).await?;
        stream.flush().await?;
        tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
        return Err(anyhow::anyhow!("No supported authentication method"));
    }

    stream.write_all(&[0x05, 0x00]).await?;

    let mut buf = [0u8; 4];
    stream.read_exact(&mut buf).await?;
    if buf[0] != 0x05 {
        return Err(anyhow::anyhow!("Invalid SOCKS5 version"));
    }
    if buf[1] != 0x01 { // CONNECT
        stream.write_all(&[0x05, 0x07, 0x00, 0x01, 0, 0, 0, 0, 0, 0]).await?;
        stream.flush().await?;
        tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
        return Err(anyhow::anyhow!("Unsupported SOCKS5 command"));
    }

    let (host, port) = match buf[3] { // ATYP
        0x01 => { // IPv4
            let mut addr_buf = [0u8; 4];
            stream.read_exact(&mut addr_buf).await?;
            let host = std::net::Ipv4Addr::from(addr_buf).to_string();
            let mut port_buf = [0u8; 2];
            stream.read_exact(&mut port_buf).await?;
            let port = u16::from_be_bytes(port_buf);
            (host, port)
        }
        0x03 => { // Domain name
            let mut len_buf = [0u8; 1];
            stream.read_exact(&mut len_buf).await?;
            let len = len_buf[0] as usize;
            let mut domain_buf = vec![0u8; len];
            stream.read_exact(&mut domain_buf).await?;
            let host = String::from_utf8(domain_buf)?;
            let mut port_buf = [0u8; 2];
            stream.read_exact(&mut port_buf).await?;
            let port = u16::from_be_bytes(port_buf);
            (host, port)
        }
        0x04 => { // IPv6
            let mut addr_buf = [0u8; 16];
            stream.read_exact(&mut addr_buf).await?;
            let host = std::net::Ipv6Addr::from(addr_buf).to_string();
            let mut port_buf = [0u8; 2];
            stream.read_exact(&mut port_buf).await?;
            let port = u16::from_be_bytes(port_buf);
            (host, port)
        }
        _ => {
            stream.write_all(&[0x05, 0x08, 0x00, 0x01, 0, 0, 0, 0, 0, 0]).await?;
            stream.flush().await?;
            tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
            return Err(anyhow::anyhow!("Unsupported address type"));
        }
    };

    // Send reply
    let reply = [0x05, 0x00, 0x00, 0x01, 0, 0, 0, 0, 0, 0];
    stream.write_all(&reply).await?;

    Ok((stream, host, port))
}

#[cfg(test)]
mod tests {
    use super::*;

    async fn setup_test_stream() -> (TcpStream, TcpStream) {
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        let client_stream = TcpStream::connect(addr).await.unwrap();
        let (server_stream, _) = listener.accept().await.unwrap();
        (client_stream, server_stream)
    }

    #[tokio::test]
    async fn test_socks5_handshake_ipv4_success() {
        let (mut client, server) = setup_test_stream().await;

        let server_handle = tokio::spawn(async move {
            handle_socks5_handshake(server).await
        });

        // Client side simulation
        client.write_all(&[0x05, 0x01, 0x00]).await.unwrap();
        let mut response = [0u8; 2];
        client.read_exact(&mut response).await.unwrap();
        assert_eq!(response, [0x05, 0x00]);

        let request = [0x05, 0x01, 0x00, 0x01, 127, 0, 0, 1, 0x00, 0x50]; // 127.0.0.1:80
        client.write_all(&request).await.unwrap();
        let mut response = [0u8; 10];
        client.read_exact(&mut response).await.unwrap();
        assert_eq!(response[0..4], [0x05, 0x00, 0x00, 0x01]);

        let result = server_handle.await.unwrap();
        assert!(result.is_ok());
        let (_, host, port) = result.unwrap();
        assert_eq!(host, "127.0.0.1");
        assert_eq!(port, 80);
    }

    #[tokio::test]
    async fn test_socks5_handshake_domain_success() {
        let (mut client, server) = setup_test_stream().await;

        let server_handle = tokio::spawn(async move {
            handle_socks5_handshake(server).await
        });

        client.write_all(&[0x05, 0x01, 0x00]).await.unwrap();
        let mut response = [0u8; 2];
        client.read_exact(&mut response).await.unwrap();
        assert_eq!(response, [0x05, 0x00]);

        let domain = b"example.com";
        let mut request = vec![0x05, 0x01, 0x00, 0x03, domain.len() as u8];
        request.extend_from_slice(domain);
        request.extend_from_slice(&[0x01, 0xBB]); // port 443
        client.write_all(&request).await.unwrap();
        let mut response = [0u8; 10];
        client.read_exact(&mut response).await.unwrap();
        assert_eq!(response[0..4], [0x05, 0x00, 0x00, 0x01]);

        let result = server_handle.await.unwrap();
        assert!(result.is_ok());
        let (_, host, port) = result.unwrap();
        assert_eq!(host, "example.com");
        assert_eq!(port, 443);
    }

    #[tokio::test]
    async fn test_unsupported_auth_method() {
        let (mut client, server) = setup_test_stream().await;

        let server_handle = tokio::spawn(async move {
            handle_socks5_handshake(server).await
        });

        client.write_all(&[0x05, 0x01, 0x02]).await.unwrap();
        let mut response = [0u8; 2];
        client.read_exact(&mut response).await.unwrap();
        assert_eq!(response, [0x05, 0xFF]);

        let result = server_handle.await.unwrap();
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_unsupported_command() {
        let (mut client, server) = setup_test_stream().await;

        let server_handle = tokio::spawn(async move {
            handle_socks5_handshake(server).await
        });

        client.write_all(&[0x05, 0x01, 0x00]).await.unwrap();
        let mut response = [0u8; 2];
        client.read_exact(&mut response).await.unwrap();

        let request = [0x05, 0x02, 0x00, 0x01, 127, 0, 0, 1, 0x00, 0x50];
        client.write_all(&request).await.unwrap();
        let mut response = [0u8; 10];
        client.read_exact(&mut response).await.unwrap();
        assert_eq!(response[1], 0x07);

        let result = server_handle.await.unwrap();
        assert!(result.is_err());
    }
}
